package com.bjfu.intelligentrecognizeserver.question_answer;

import com.alibaba.fastjson.JSON;
import com.bjfu.intelligentrecognizeserver.Util.FileTransferDataStore;
import com.bjfu.intelligentrecognizeserver.Util.SessionImageStore;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.zhipu.oapi.ClientV4;
import com.zhipu.oapi.Constants;
import com.zhipu.oapi.service.v4.model.*;
import io.reactivex.Flowable;
import jakarta.websocket.server.ServerEndpoint;
import org.jetbrains.annotations.NotNull;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.*;

import java.io.IOException;
import java.util.*;
import java.util.concurrent.CopyOnWriteArraySet;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.atomic.AtomicBoolean;

import static com.bjfu.intelligentrecognizeserver.Util.FileUtil.imageToBase64;
import static com.zhipu.oapi.service.v4.deserialize.MessageDeserializeFactory.defaultObjectMapper;

//@ServerEndpoint(value = "/questionAnswer/{userId}")
//@Component
public class QuestionAnswerServer implements WebSocketHandler {
    private static final Logger log = LoggerFactory.getLogger(QuestionAnswerServer.class);

    //存放每个客户端对应的WebSocket对象
    private static final CopyOnWriteArraySet<QuestionAnswerServer> webSocketSet = new CopyOnWriteArraySet<>();
    private WebSocketSession mSession;
    private String userId;

    //SDK Client
    private static final ClientV4 client;
    private static final ExecutorService threadPool = Executors.newCachedThreadPool();
    private static final ObjectMapper mapper = defaultObjectMapper();
    private final String requestIdTemplate = "intelligent-recognize-server-%s";

    static {
        //初始化大模型SDK客户端
        client = new ClientV4.Builder(
                "fd600744763eb6a42236eab6a15cd904.6gqfm0z4bThwpYzw"
        ).build();
    }

    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        log.info("websocket connection established");
        log.info("session id:{} ", session.getId());
        String uri = Objects.requireNonNull(session.getUri()).toString();
        this.userId = uri.substring(uri.lastIndexOf("/") + 1);
        log.info("userId:{} ", this.userId);
        log.info("websocket connection established for userId: {}", userId);
        this.mSession = session;
        webSocketSet.add(this);
    }

    @Override
    public void handleMessage(@NotNull WebSocketSession session, WebSocketMessage<?> message) throws Exception {
        //此处应该把message传给大模型，大模型回答得到答案，返回给服务器，服务器再将答案给客户端
        String clientMessage = (String) message.getPayload();
        log.info("received message:{}", clientMessage);

        //使用线程池处理每一个请求
        threadPool.submit(() -> {
            //handleClientMessage(clientMessage, session);
            //handleSpeciesRecognitionMessage(clientMessage, session);
            //testImageToWord();
            ObjectMapper objectMapper = new ObjectMapper();
            Map<String, Object> messageMap;
            try {
                messageMap = objectMapper.readValue(clientMessage, Map.class);
            } catch (JsonProcessingException e) {
                log.error("消息解析失败：{}", e.getMessage());
                try {
                    session.sendMessage(new TextMessage("消息格式错误"));
                } catch (IOException ex) {
                    throw new RuntimeException(ex);
                }
                throw new RuntimeException(e);
            }

            String text = (String) messageMap.get("text");
            String imageBase64 = (String) messageMap.get("image");

            if (imageBase64 != null && !imageBase64.isEmpty()) {
                SessionImageStore.setImage(userId, imageBase64);
                handleImageAndTextMessage(text, imageBase64, session);
            } else {
                //仅文字消息，检查是否有关联的图片
                String image = SessionImageStore.getImage(userId);
                if (image != null && !image.isEmpty()) {
                    handleTextWithImageMessage(text, image, session);
                } else {
                    handleTextMessage(text, session);
                }
            }

        });

    }

    private void handleTextMessage(String text, WebSocketSession session) {
        List<ChatMessage> messages = new ArrayList<>();

        // 仅添加文本消息
        ChatMessage chatMessage = new ChatMessage(ChatMessageRole.USER.value(), text);
        messages.add(chatMessage);

        String requestId = String.format(requestIdTemplate, System.currentTimeMillis());
        ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
                .stream(Boolean.TRUE)
                .model("glm-4v-flash")
                .invokeMethod(Constants.invokeMethodAsync)
                .messages(messages)
                .requestId(requestId)
                .build();
        sseResponse(session, chatCompletionRequest);
    }

    private void handleTextWithImageMessage(String text, String image, WebSocketSession session) {
        List<ChatMessage> messages = new ArrayList<>();

        // 添加文本消息
        ChatMessage chatMessage = new ChatMessage(ChatMessageRole.USER.value(), text);
        messages.add(chatMessage);

        // 添加关联的图片消息
        ChatMessage imageMessage = new ChatMessage(ChatMessageRole.USER.value(), Collections.singletonList(
                Map.of(
                        "type", "image_url",
                        "image_url", Map.of("url", image)
                )
        ));
        messages.add(imageMessage);

        String requestId = String.format(requestIdTemplate, System.currentTimeMillis());
        ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
                .stream(Boolean.TRUE)
                .model("glm-4v-flash")
                .invokeMethod(Constants.invokeMethodAsync)
                .messages(messages)
                .requestId(requestId)
                .build();
        sseResponse(session, chatCompletionRequest);
    }

    private void handleImageAndTextMessage(String text, String imageBase64, WebSocketSession session) {
        List<ChatMessage> messages = new ArrayList<>();

        // 添加文本消息
        ChatMessage chatMessage = new ChatMessage(ChatMessageRole.USER.value(), text);
        messages.add(chatMessage);

        // 添加图片消息
        ChatMessage imageMessage = new ChatMessage(ChatMessageRole.USER.value(), Collections.singletonList(
                Map.of(
                        "type", "image_url",
                        "image_url", Map.of("url", imageBase64)
                )
        ));
        messages.add(imageMessage);

        String requestId = String.format(requestIdTemplate, System.currentTimeMillis());
        ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
                .stream(Boolean.TRUE)
                .model("glm-4v-flash")
                .invokeMethod(Constants.invokeMethodAsync)
                .messages(messages)
                .requestId(requestId)
                .build();
        sseResponse(session, chatCompletionRequest);
    }

    private void handleClientMessage(String clientMessage, WebSocketSession session) {
        ChatMessage chatMessage = new ChatMessage(ChatMessageRole.USER.value(), clientMessage);
        List<ChatMessage> messages = new ArrayList<>();
        messages.add(chatMessage);
        String requestId = String.format(requestIdTemplate, System.currentTimeMillis());
        ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
                .stream(Boolean.TRUE)
                .model(Constants.ModelChatGLM4V)
                .invokeMethod(Constants.invokeMethodAsync)
                .messages(messages)
                .requestId(requestId)
                .build();
        sseResponse(session, chatCompletionRequest);

    }

    private void handleSpeciesRecognitionMessage(String clientMessage, WebSocketSession session) {
        List<ChatMessage> messages = getChatMessages(clientMessage);
        String requestId = String.format(requestIdTemplate, System.currentTimeMillis());
        ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
                .stream(Boolean.TRUE)
                .model("glm-4v-flash")
                .invokeMethod(Constants.invokeMethodAsync)
                .messages(messages)
                .requestId(requestId)
                .build();
        sseResponse(session, chatCompletionRequest);

    }

    private void sseResponse(WebSocketSession session, ChatCompletionRequest chatCompletionRequest) {
        ModelApiResponse sseModelApiResp = client.invokeModelApi(chatCompletionRequest);
        if (sseModelApiResp.isSuccess()) {
            AtomicBoolean isFirst = new AtomicBoolean(true);
            ChatMessageAccumulator chatMessageAccumulator = mapStreamToAccumulator(sseModelApiResp.getFlowable())
                    .doOnNext(accumulator -> {
                        {
                            if (isFirst.getAndSet(false)) {
                                session.sendMessage(new TextMessage("Response: "));
                            }
                            if (accumulator.getDelta() != null && accumulator.getDelta().getTool_calls() != null) {
                                String jsonString = mapper.writeValueAsString(accumulator.getDelta().getTool_calls());
                                System.out.println("tool_calls: " + jsonString);
                            }
                            if (accumulator.getDelta() != null && accumulator.getDelta().getContent() != null) {
                                session.sendMessage(new TextMessage(accumulator.getDelta().getContent()));
                            }
                        }
                    })
                    .doOnComplete(() -> {
                        session.sendMessage(new TextMessage("Response end"));
                    })
                    .lastElement()
                    .blockingGet();
            Choice choice = new Choice();
            List<Choice> choices = new ArrayList<>();
            choices.add(choice);
            ModelData data = new ModelData();
            data.setChoices(choices);
            data.setUsage(chatMessageAccumulator.getUsage());
            data.setId(chatMessageAccumulator.getId());
            data.setCreated(chatMessageAccumulator.getCreated());
            data.setRequestId(chatCompletionRequest.getRequestId());
            sseModelApiResp.setFlowable(null);
            sseModelApiResp.setData(data);
        }
        log.info(JSON.toJSONString(sseModelApiResp));
    }

    private @NotNull List<ChatMessage> getChatMessages(String clientMessage) {
        List<ChatMessage> messages = new ArrayList<>();
        List<Map<String, Object>> contentList = new ArrayList<>();
        Map<String, Object> textMap = new HashMap<>();
        textMap.put("type", "text");
        textMap.put("text", clientMessage);
        Map<String, Object> typeMap = new HashMap<>();
        typeMap.put("type", "image_url");
        Map<String, Object> urlMap = new HashMap<>();
        String image = imageToBase64(FileTransferDataStore.filePath.get(FileTransferDataStore.filePath.size() - 1));
        log.info(FileTransferDataStore.filePath.get(0));
        urlMap.put("url", image);
        typeMap.put("image_url", urlMap);
        contentList.add(textMap);
        contentList.add(typeMap);
        ChatMessage chatMessage = new ChatMessage(ChatMessageRole.USER.value(), contentList);
        messages.add(chatMessage);
        return messages;
    }


    private Flowable<ChatMessageAccumulator> mapStreamToAccumulator(Flowable<ModelData> flowable) {
        return flowable.map(chunk -> new ChatMessageAccumulator(chunk.getChoices().get(0).getDelta(), null,
                chunk.getChoices().get(0), chunk.getUsage(), chunk.getCreated(), chunk.getId()));
    }


    @Override
    public void handleTransportError(WebSocketSession session, Throwable exception) throws Exception {
        log.info("websocket连接异常,session id: {}", session.getId());
        log.error(exception.getMessage());
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, @NotNull CloseStatus closeStatus) throws Exception {
        log.info("websocket 连接关闭,session id: {}", session.getId());
        webSocketSet.remove(this);
        //移除关联图片
        SessionImageStore.removeImage(userId);
    }

    //是否消息一段一段发送
    @Override
    public boolean supportsPartialMessages() {
        return false;
    }
}
